from base.base_trainer import BaseTrain
from utils.callbacks import distance_metrics, categorical_metrics, distance_mixed_metrics, categorical_mixed_metrics

from keras.callbacks import ModelCheckpoint, TensorBoard, EarlyStopping

import os


class SiameseModelTrainer(BaseTrain):
    def __init__(self, model, train_data, valid_data, config):
        super(SiameseModelTrainer, self).__init__(model, train_data, valid_data, config=config)
        self.callbacks = []
        self.loss = []
        self.acc = []
        self.val_loss = []
        self.val_acc = []
        self.val_precision = []
        self.val_recall = []
        self.val_f1 = []
        self.val_auc = []
        self.init_callbacks()

    def init_callbacks(self):
        if self.config.exp_name in ('siamese_categorical_word', 'siamese_categorical_char',
                                    'siamese_categorical_mixed'):
            self.callbacks.append(categorical_metrics)
        else:
            self.callbacks.append(distance_metrics)

        self.callbacks.append(
            ModelCheckpoint(
                filepath=os.path.join(self.config.checkpoint_dir, '%s.hdf5' % self.config.exp_name),
                monitor=self.config.checkpoint_monitor,
                mode=self.config.checkpoint_mode,
                save_best_only=self.config.checkpoint_save_best_only,
                save_weights_only=self.config.checkpoint_save_weights_only,
                verbose=self.config.checkpoint_verbose,
            )
        )

        self.callbacks.append(
            TensorBoard(
                log_dir=self.config.tensorboard_log_dir,
                write_graph=self.config.tensorboard_write_graph,
            )
        )

        self.callbacks.append(
            EarlyStopping(
                monitor=self.config.early_stopping_monitor,
                patience=self.config.early_stopping_patience,
                mode=self.config.early_stopping_mode
            )
        )

    def train(self):
        history = self.model.fit(
            [self.train_data_word[0], self.train_data_word[1]], self.train_data_word[2],
            epochs=self.config.num_epochs,
            verbose=self.config.verbose_training,
            batch_size=self.config.batch_size,
            validation_data=([self.valid_data_word[0], self.valid_data_word[1]], self.valid_data_word[2]),
            callbacks=self.callbacks
        )
        # if self.config.distance == 'cosine':
        #     acc_name = 'cosine_acc'
        # elif self.config.distance == 'eucl':
        #     acc_name = 'eucl_acc'
        # else:
        #     acc_name = 'acc'
        self.loss.extend(history.history['loss'])
        self.acc.extend(history.history['acc'])
        self.val_loss.extend(history.history['val_loss'])
        self.val_acc.extend(history.history['val_acc'])
        # self.val_precision.extend(history.history['val_precision'])
        # self.val_recall.extend(history.history['val_recall'])
        # self.val_f1.extend(history.history['val_f1'])
        # self.val_auc.extend(history.history['val_auc'])

        # self.model.save(os.path.join(self.config.checkpoint_dir, '%s.hdf5' % self.config.exp_name))


class SiameseMixedTrainer(BaseTrain):
    def __init__(self, model, train_data_word, valid_data_word, train_data_char, valid_data_char, config):
        super(SiameseMixedTrainer, self).__init__(model,
                                                  train_data_word, valid_data_word,
                                                  train_data_char, valid_data_char,
                                                  config)
        self.callbacks = []
        self.loss = []
        self.acc = []
        self.val_loss = []
        self.val_acc = []
        self.val_precision = []
        self.val_recall = []
        self.val_f1 = []
        self.val_auc = []
        self.init_callbacks()

    def init_callbacks(self):
        if self.config.exp_name in ('siamese_categorical_word', 'siamese_categorical_char',
                                    'siamese_categorical_mixed', 'ensemble'):
            self.callbacks.append(categorical_mixed_metrics)
        else:
            self.callbacks.append(distance_mixed_metrics)

        self.callbacks.append(
            ModelCheckpoint(
                filepath=os.path.join(self.config.checkpoint_dir, '%s.hdf5' % self.config.exp_name),
                monitor=self.config.checkpoint_monitor,
                mode=self.config.checkpoint_mode,
                save_best_only=self.config.checkpoint_save_best_only,
                save_weights_only=self.config.checkpoint_save_weights_only,
                verbose=self.config.checkpoint_verbose,
            )
        )

        self.callbacks.append(
            TensorBoard(
                log_dir=self.config.tensorboard_log_dir,
                write_graph=self.config.tensorboard_write_graph,
            )
        )

        self.callbacks.append(
            EarlyStopping(
                monitor=self.config.early_stopping_monitor,
                patience=self.config.early_stopping_patience,
                mode=self.config.early_stopping_mode
            )
        )

    def train(self):
        history = self.model.fit(
            [self.train_data_word[0], self.train_data_word[1], self.train_data_char[0], self.train_data_char[1]],
            self.train_data_word[2],
            epochs=self.config.num_epochs,
            verbose=self.config.verbose_training,
            batch_size=self.config.batch_size,
            validation_data=([self.valid_data_word[0], self.valid_data_word[1],
                              self.valid_data_char[0], self.valid_data_char[1]],
                             self.valid_data_word[2]),
            callbacks=self.callbacks,
        )
        # if self.config.distance == 'cosine':
        #     acc_name = 'cosine_acc'
        # elif self.config.distance == 'eucl':
        #     acc_name = 'eucl_acc'
        # else:
        #     acc_name = 'acc'
        self.loss.extend(history.history['loss'])
        self.acc.extend(history.history['acc'])
        self.val_loss.extend(history.history['val_loss'])
        self.val_acc.extend(history.history['val_acc'])
        # self.val_precision.extend(history.history['val_precision'])
        # self.val_recall.extend(history.history['val_recall'])
        # self.val_f1.extend(history.history['val_f1'])
        # self.val_auc.extend(history.history['val_auc'])

        # self.model.save(os.path.join(self.config.checkpoint_dir, '%s.hdf5' % self.config.exp_name))


class EnsembleTrainer(object):
    def __init__(self, model, x, y):
        self.model = model
        self.x = x
        self.y = y

    def train(self):
        self.model.fit(self.x, self.y)

